import pandas as pd
import numpy as np
#%% Country replacement
def get_country_name():
    replacement = {'United Kingdom of Great Britain and Northern Ireland': 'United Kingdom of Great Britain and Northern Ireland (the)',
                   'United States of America': 'United States of America (the)',
                   'USA': 'United States of America (the)',
                   'United Arab Emirates': 'United Arab Emirates (the)',
                   'Republic of Korea': 'Korea (the Republic of)',
                   'Taiwan, Greater China': 'Taiwan (Province of China)',
                   'Taiwan, China': 'Taiwan (Province of China)',
                   'China, Hong Kong Special Administrative Region': 'Hong Kong',
                   'Russian Federation': 'Russian Federation (the)',
                   'Cayman Islands': 'Cayman Islands (the)',
                   'Netherlands': 'Netherlands (the)',
                   'Marshall Islands': 'Marshall Islands (the)',
                   'Philippines': 'Philippines (the)',
                   ####
                   'United Kingdom': 'United Kingdom of Great Britain and Northern Ireland (the)',
                   'South Korea': 'Korea (the Republic of)',
                   'Taiwan': 'Taiwan (Province of China)',
                   'Russia': 'Russian Federation (the)',
                   'Czech Republic': 'Czechia',
                   'Dominican Republic': 'Dominican Republic (the)',
                   'Faroe Islands': 'Faroe Islands (the)',
                   'British Virgin Islands': 'Virgin Islands (British)',
                   'China, Macao Special Administrative Region': 'Macao',
                   'Cocos (Keeling) Islands': 'Cocos (Keeling) Islands (the)',
                   'Bosnia & Herzegovina': 'Bosnia and Herzegovina',
                   'North Macedonia': 'Republic of North Macedonia',
                   'United States Minor Outlying Islands': 'United States Minor Outlying Islands (the)'
                   }
    return replacement
#%% 
def get_SAR():
    ARTYPE = 'SAR'
    SAR100 = [[ 'CO2-'+ARTYPE+'-100', 1], ['CH4-'+ARTYPE+'-100', 21], ['N2O-'+ARTYPE+'-100',  310], ['HFCs-'+ARTYPE+'-100', 11700], ['SF6-'+ARTYPE+'-100', 23900], ['NF3-'+ARTYPE+'-100', np.nan] ]
    SAR20 = [[ 'CO2-'+ARTYPE+'-20', 1], ['CH4-'+ARTYPE+'-20', 56], ['N2O-'+ARTYPE+'-20', 280], ['HFCs-'+ARTYPE+'-20', 9100], ['SF6-'+ARTYPE+'-20', 16300], ['NF3-'+ARTYPE+'-20',np.nan] ]
    SAR = pd.concat((pd.DataFrame(SAR100), pd.DataFrame(SAR20)))
    SAR.columns = ['mrg', 'ConversionFactor']
    return SAR

def get_TAR():
    ARTYPE = 'TAR'
    TAR100 = [[ 'CO2-'+ARTYPE+'-100', 1], ['CH4-'+ARTYPE+'-100', 23], ['N2O-'+ARTYPE+'-100', 296], ['HFCs-'+ARTYPE+'-100', 120], ['SF6-'+ARTYPE+'-100', 22200], ['NF3-'+ARTYPE+'-100', np.nan] ]
    TAR20 = [[ 'CO2-'+ARTYPE+'-20', 1], ['CH4-'+ARTYPE+'-20', 62], ['N2O-'+ARTYPE+'-20', 275], ['HFCs-'+ARTYPE+'-20', 410], ['SF6-'+ARTYPE+'-20', 15100], ['NF3-'+ARTYPE+'-20',np.nan] ]
    TAR = pd.concat((pd.DataFrame(TAR100), pd.DataFrame(TAR20)))
    TAR.columns = ['mrg', 'ConversionFactor']
    return TAR
def get_AR4():
    AR4100 = [[ 'CO2-AR4-100', 1], ['CH4-AR4-100', 25], ['N2O-AR4-100', 298], ['HFCs-AR4-100', 124], ['SF6-AR4-100', 22800], ['NF3-AR4-100', 17200] ]
    AR420 = [[ 'CO2-AR4-20', 1], ['CH4-AR4-20', 72], ['N2O-AR4-20', 289], ['HFCs-AR4-20', 437], ['SF6-AR4-20', 16300], ['NF3-AR4-20', 12300] ]
    AR4 = pd.concat((pd.DataFrame(AR4100), pd.DataFrame(AR420)))
    AR4.columns = ['mrg', 'ConversionFactor']
    return AR4

def get_AR5():
    ARTYPE = 'AR5'
    AR5100 = [[ 'CO2-'+ARTYPE+'-100', 1], ['CH4-'+ARTYPE+'-100', 28], ['N2O-'+ARTYPE+'-100', 265], ['HFCs-'+ARTYPE+'-100', 138], ['SF6-'+ARTYPE+'-100', 23500], ['NF3-'+ARTYPE+'-100', 16100] ]
    AR520 = [[ 'CO2-'+ARTYPE+'-20', 1], ['CH4-'+ARTYPE+'-20', 84], ['N2O-'+ARTYPE+'-20', 264], ['HFCs-'+ARTYPE+'-20', 506], ['SF6-'+ARTYPE+'-20', 17500], ['NF3-'+ARTYPE+'-20', 12800] ]
    AR5 = pd.concat((pd.DataFrame(AR5100), pd.DataFrame(AR520)))
    AR5.columns = ['mrg', 'ConversionFactor']
    return AR5

def get_AR6(latest_methane=28.4):
    AR6100 = [[ 'CO2-AR6-100', 1], ['CH4-AR6-100', 28.4], ['N2O-AR6-100', 273], ['HFCs-AR6-100', 164], ['SF6-AR6-100', 25200], ['NF3-AR6-100', 17400] ]
    AR620 = [[ 'CO2-AR6-20', 1], ['CH4-AR6-20', 81.1], ['N2O-AR6-20', 273], ['HFCs-AR6-20', 591], ['SF6-AR6-20', 18300], ['NF3-AR6-20', 13400] ]
    AR6 = pd.concat((pd.DataFrame(AR6100), pd.DataFrame(AR620)))
    AR6.columns = ['mrg', 'ConversionFactor']
    return AR6

def get_latest_100(ARTYPE, latest_methane=28.4):
    if ARTYPE == 'AR5':
        ARL = [[ 'CO2-100', 1], ['CH4-100', 28], ['N2O-100', 265], ['HFCs-100', 138], ['SF6-100', 23500], ['NF3-100', 16100] ]
    if ARTYPE == 'AR6':
        ARL = [[ 'CO2-100', 1], ['CH4-100', latest_methane], ['N2O-100', 273], ['HFCs-100', 164], ['SF6-100', 25200], ['NF3-100', 17400] ]
        
    ARL = pd.DataFrame(ARL, columns = ['mrg', 'EqualConversionFactor'])
    return ARL
def get_latest_20(ARTYPE, latest_methane=81.1):
    if ARTYPE == 'AR5':
        ARL = [[ 'CO2-20', 1], ['CH4-20', 84], ['N2O-20', 264], ['HFCs-20', 506], ['SF6-20', 17500], ['NF3-20', 12800] ]
    if ARTYPE == 'AR6':
        ARL = [[ 'CO2-20', 1], ['CH4-20', latest_methane], ['N2O-20', 273], ['HFCs-20', 591], ['SF6-20', 18300], ['NF3-20', 13400] ]
        
    ARL = pd.DataFrame(ARL, columns = ['mrg', 'EqualConversionFactor'])
    return ARL



#==================================================
#=== This is the most important function of all ===
#==================================================
def Counterfactual(D, Z, YEARS, latest_AsRp,latest_methane=28.4, horizon = 100):
    #=== Only focus on the years with common latest GWP
    X = D[D.ReportingYear.isin(YEARS)]
    X = X.drop(columns = ['mrg'])
    X.insert(0,'mrg', (X['Gas']+'-'+X['GWP_adj']).values.tolist())
    #=== Only focus on those observation for which you can recover a GWP
    X = X[X.GWP_adj.isin(['SAR-20', 'SAR-100', 'TAR-20', 'TAR-100', 'AR4-20', 'AR4-100', 'AR5-20', 'AR5-100', 'AR6-20', 'AR6-100'])]
    X = X.merge(Z)
    #=== Convert tCO2e into Natural unit (the conversion factor is the GWP)
    X['Value (NatUnit)'] = X['Value (tCO2e)'].astype(float)/X['ConversionFactor'].astype(float)

    if horizon == 100:
        late_AR = get_latest_100(latest_AsRp,latest_methane)
    elif horizon == 20:
        late_AR = get_latest_20(latest_AsRp,latest_methane)

    #=== Create a column to merge the Gas type with the counterfactual GWP
    X['mrg'] = X['Gas']+'-'+str(horizon)
    X = X.merge(late_AR)
    #==== Convert the natural unit gas into GWP (This is the counterfactual)
    X['Value (Equalised)'] = X['Value (NatUnit)'].astype(float)*X['EqualConversionFactor'].astype(float)
    X['Delta'] = X['Value (tCO2e)'].astype(float) - X['Value (Equalised)'].astype(float)

    #== The following is basically the ratio of GWP
    X['Delta,%'] = X['Value (tCO2e)'].astype(float)/X['Value (Equalised)'].astype(float)

    return X